import numpy as np
import argparse
import cv2

from paddle.inference import Config
from paddle.inference import create_predictor
from paddle.inference import PrecisionType


# this is a simple resnet block for dynamci test.
def init_predictor(args):
    if args.model_dir is not "":
      config = Config(args.model_dir)
    else:
      config = Config(args.model_file, args.params_file)
    config.enable_memory_optim()
    config.enable_use_gpu(100, 0)

    # using dynamic shpae mode, the max_batch_size will be ignored.
    config.enable_tensorrt_engine(workspace_size=1 << 30,
                                  max_batch_size=1,
                                  min_subgraph_size=5,
                                  precision_mode=PrecisionType.Float32,
                                  use_static=False,
                                  use_calib_mode=False)

    head_number = 12

    names = [
        "placeholder_0", "placeholder_1", "placeholder_2", "stack_0.tmp_0"
    ]
    min_input_shape = [1, 1, 1]
    max_input_shape = [100, 128, 1]
    opt_input_shape = [10, 60, 1]

    config.set_trt_dynamic_shape_info(
        {
            names[0]: min_input_shape,
            names[1]: min_input_shape,
            names[2]: min_input_shape,
            names[3]: [1, head_number, 1, 1]
        }, {
            names[0]: max_input_shape,
            names[1]: max_input_shape,
            names[2]: max_input_shape,
            names[3]: [100, head_number, 128, 128]
        }, {
            names[0]: opt_input_shape,
            names[1]: opt_input_shape,
            names[2]: opt_input_shape,
            names[3]: [10, head_number, 60, 60]
        })
    predictor = create_predictor(config)
    return predictor


def run(predictor, data):
    # copy data to input tensor
    input_names = predictor.get_input_names()
    for i, name in enumerate(input_names):
        input_tensor = predictor.get_input_handle(name)
        input_tensor.reshape(data[i].shape)
        input_tensor.copy_from_cpu(data[i].copy())

    # do the inference
    predictor.run()

    results = []
    # get out data from output tensor
    output_names = predictor.get_output_names()
    for i, name in enumerate(output_names):
        output_tensor = predictor.get_output_handle(name)
        output_data = output_tensor.copy_to_cpu()
        results.append(output_data)

    return results


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model_file",
        type=str,
        default="",
        help="Model filename, Specify this when your model is a combined model."
    )
    parser.add_argument(
        "--params_file",
        type=str,
        default="",
        help=
        "Parameter filename, Specify this when your model is a combined model."
    )
    parser.add_argument(
        "--model_dir",
        type=str,
        default="",
        help=
        "Model dir, If you load a non-combined model, specify the directory of the model."
    )
    parser.add_argument("--use_gpu",
                        type=int,
                        default=0,
                        help="Whether use gpu.")
    return parser.parse_args()


if __name__ == '__main__':
    args = parse_args()
    pred = init_predictor(args)
    in1 = np.ones((1, 128, 1)).astype(np.int64)
    in2 = np.ones((1, 128, 1)).astype(np.int64)
    in3 = np.ones((1, 128, 1)).astype(np.int64)
    in4 = np.ones((1, 128, 1)).astype(np.float32)
    result = run(pred, [in1, in2, in3, in4])
    print(result)
